package optimization;

import static optimization.LinearAlgebra.*;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import optimization.OldLogger;

import common.statistics.Series;


public class NelderMeadStateMachine {

	/**
	 * Function value at a point
	 * @author risto
	 *
	 */
	class PointValue implements Comparable<PointValue> {
		double[] x;
		double fx;
		public PointValue(double[] p) {
			this.x = p;
		}
		@Override
		public int compareTo(PointValue o) {
			return Double.compare(fx, o.fx);
		}
	}

	/**
	 * State machine state
	 * @author risto
	 *
	 */
	private abstract class State {
		public final String name;
		public State(String name) {
			this.name = name;
			stateMap.put(name, this);
		}
		public abstract double[] getX();
		public abstract void modelComputed(double value);
	}


	/**
	 * Simplex points
	 */
	List<PointValue> points;
	
	/**
	 * Best, second-worst, worst
	 */
	double fl, fs, fh; 

	/**
	 * centroid
	 */
	double[] c;
	
	/**
	 * reflection point
	 */
	double[] xr;
	double fr;

	
	/**
	 * Constants
	 */
	protected double alpha = 1, beta = 0.5, gamma = 2, delta = 0.5;
	
	/**
	 * Function (value)
	 */
	Function g;

	/**
	 * Constraints mapping
	 */
	VectorFunction cons;
	
	/**
	 * Dennis-Woods termination threshold
	 */
	double nu = 0.0001;

	boolean done = false;
	
	// samples
	int N=1;
	// random seed
//	long seed=0;

	// current iteration
	int iter = 0;
	// max iterations
	int MAX_ITER = 10; // 5000;

	// y value series
	Series ysn = new Series();
	
	// current state
	State state;

	// map of all states
	Map<String, State> stateMap = new HashMap<String, NelderMeadStateMachine.State>();

	
	StringBuffer path = new StringBuffer();
	
	
	public NelderMeadStateMachine(Function g, VectorFunction cons) {
		super();
		this.g = g;
		this.cons = cons;
	}
	
	public NelderMeadStateMachine(Function g, VectorFunction cons, double terminationThreshold) {
		super();
		this.g = g;
		this.cons = cons;
		this.nu = terminationThreshold;
	}
	
	public void setStochasticParameters(int N)//, long seed)
	{
		this.N = N;
//		this.seed = seed;
	}
/*
	public long getRandSeed()
	{
		return seed;
	}
	*/	
	
	protected boolean dennisWoodsTerminationTest()
	{
		double normX0 = Math.max(1,vNorm2(points.get(0).x));
		double maxNorm = Double.MIN_VALUE;
		for (int i = 1; i < 3; i++) {
			double norm = vNorm2(vSub(points.get(0).x, points.get(i).x));
			if (norm > maxNorm) maxNorm = norm;
		}
		if (maxNorm / normX0 < nu) return true;
		return false;
	}
	
	public void initSearch(double[] x0i, double[] h0)
	{
		points = new ArrayList<NelderMeadStateMachine.PointValue>();
		
		points.add(new PointValue(x0i));
		points.add(new PointValue(vAdd(x0i, vSelect(h0, 0))));
		points.add(new PointValue(vAdd(x0i, vSelect(h0, 1))));
		
		iter = 0;
		path = new StringBuffer();
		done = false;
		goTo("INIT");
	}

	public boolean isDone()
	{
		return done;
	}


	
	/**
	 * Go to next state
	 * @param nextState state name
	 */
	protected void goTo(String nextState)
	{
		state = stateMap.get(nextState);
	}

	/**
	 * Initial state, compute values at {x_i}
	 */
	protected final State INIT = new State("INIT") {

		int initVar = 0;

		@Override
		public double[] getX() {
			return points.get(initVar).x;
		}

		@Override
		public void modelComputed(double value) {
			points.get(initVar).fx = value;
			initVar++;
			if (initVar == 3) {
				startIteration();
				initVar = 0; // for next round
			}
		}
	};

	/**
	 * Compute reflection point
	 */
	protected final State REFLECT = new State("REFLECT") {
		@Override
		public double[] getX() {
			xr = cons.evaluate(vAdd(c, vMul(getAlpha(), vSub(c, points.get(2).x))));
			return xr;
		}

		@Override
		public void modelComputed(double value) {
			fr = value;
			if (fl <= fr && fr < fs) { // Reflect 
				path.append(".RR");
				points.get(2).x = xr;
				points.get(2).fx = fr;
				nextIteration(); 
			}
			// expand ?
			else if (fr < fl) { goTo("EXPAND");	}
			// contr outside
			else if (fr >= fs && fs <= fr && fr < fh) { goTo("CONTRACT_OUTSIDE"); }
			// contr inside
			else if (fr >= fs && fr >= fh) { goTo("CONTRACT_INSIDE"); }
			// shrink
			else { goTo("SHRINK_1"); }
		}
	};

	/**
	 * Expand simplex
	 */
	protected final State EXPAND = new State("EXPAND") {
		double[] xe;
		double fe;

		@Override
		public double[] getX() {
			xe = cons.evaluate(vAdd(c, vMul(getGamma(), vSub(xr, c))));
			return xe;
		}

		@Override
		public void modelComputed(double value) {
			fe = value;
			if (fe < fr) {
				path.append(".EE");
				points.get(2).x = xe;
				points.get(2).fx = fe;
			} else {
				path.append(".ER");
				points.get(2).x = xr;
				points.get(2).fx = fr;
			}
			nextIteration();
		}
	};

	/**
	 * Contract outside
	 */
	protected final State CONTRACT_OUTSIDE = new State("CONTRACT_OUTSIDE") {
		double[] xc;
		double fc;

		@Override
		public double[] getX() {
			xc = cons.evaluate(vAdd(c, vMul(getBeta(), vSub(xr, c))));
			return xc;
		}

		@Override
		public void modelComputed(double value) {
			fc = value;
			if (fc <= fr) {
				path.append(".CO");
				points.get(2).x = xc;
				points.get(2).fx = fc;
				nextIteration();
			}
			else {
				goTo("SHRINK_1");
			}
		}
	};

	/**
	 * Contract inside
	 */
	protected final State CONTRACT_INSIDE = new State("CONTRACT_INSIDE") {
		double[] xc;
		double fc;

		@Override
		public double[] getX() {
			double[] xh = points.get(2).x;
			xc = cons.evaluate(vAdd(c, vMul(getBeta(), vSub(xh, c))));
			return xc;
		}

		@Override
		public void modelComputed(double value) {
			fc = value;
			if (fc < fh) {
				path.append(".CI");
				points.get(2).x = xc;
				points.get(2).fx = fc;
				nextIteration();
			}
			else {
				goTo("SHRINK_1");
			}
		}
	};

	/**
	 * Shrink (point 1)
	 */
	protected final State SHRINK_1 = new State("SHRINK_1") {
		double[] xs1;
		double fs1;
		
		@Override
		public double[] getX() {
			xs1 = cons.evaluate(vAdd(points.get(0).x, vMul(getDelta(), vSub(points.get(1).x, points.get(0).x))));
			return xs1;
		}

		@Override
		public void modelComputed(double value) {
			path.append(".S1");
			fs1 = value;
			points.get(1).x = xs1;
			points.get(1).fx = fs1;
			goTo("SHRINK_2");
		}
	};

	/**
	 * Shrink (point 2)
	 */
	protected final State SHRINK_2 = new State("SHRINK_2") {
		double[] xs2;
		double fs2;
		
		@Override
		public double[] getX() {
			xs2 = cons.evaluate(vAdd(points.get(0).x, vMul(getDelta(), vSub(points.get(2).x, points.get(0).x))));
			return xs2;
		}

		@Override
		public void modelComputed(double value) {
			path.append(".S2");
			fs2 = value;
			points.get(2).x = xs2;
			points.get(2).fx = fs2;
			nextIteration();
		}
	};

	/**
	 * End state
	 */
	protected final State END = new State("END") {
		@Override
		public double[] getX() {
			return getRoot();
		}

		@Override
		public void modelComputed(double value) {
		}
	};
	

	/**
	 * Get current x
	 * @return
	 */
	public double[] getX()
	{
		return state.getX();
	}
	
	/**
	 * 
	 * @param ys
	 */
	public void modelComputed(Series ys)
	{
		// merge new observations with current data vector
		ysn.merge(ys);
		
		OldLogger.format("YSN: N=%d, m=%.4f, s=%.4f", ysn.count(), ysn.mean(), ysn.stdev());
		
		// if less than required, don't process
		if (ysn.count() < N) return;

		// N or more observations, process 
		state.modelComputed(g.evaluate(ysn.meanAndDev()));
		
		// clear series, data used
		ysn.clear();
	}
	

	/**
	 * get root estimate
	 * @return
	 */
	public double[] getRoot()
	{
		return points.get(0).x;
	}

	/**
	 * Start of iteration loop
	 */
	public void startIteration()
	{		
		// sort points by value
		Collections.sort(points);
		fl = points.get(0).fx; // best        0
		fs = points.get(1).fx; // 2nd worst   n-2
		fh = points.get(2).fx; // worst       n-1

		// termination test(s)
//		if (dxTerminationTest()) return points.get(0).x;
		if (dennisWoodsTerminationTest()) {
			OldLogger.format("NM--  path: "  + path.toString());
			done = true; return; 
		}
		
		OldLogger.format("NM%d   %s, %s, %s     %.5f, %.5f, %.5f ",
				iter, 
				vPrint(points.get(0).x), vPrint(points.get(1).x), vPrint(points.get(2).x),
				fl, fs, fh);

		// compute centroid
		c = vAdd(points.get(0).x, points.get(1).x);
		c = vMul(1.0/2.0, c);

		
		// compute reflection point
// 		fr = f.evaluate(xr, N, seed);
		goTo("REFLECT");

//		boolean tmp = reflect() || expand() || contract();		
	}
	
	private void nextIteration()
	{
		// update iteration count
		iter++;
		if (iter >= MAX_ITER) { goTo("END"); done = true; return; }
		
		// next iteration
		startIteration();
	}
		

	protected double getAlpha() {
		return alpha;
	}

	protected double getBeta() {
		return beta;
	}

	protected double getGamma() {
		return gamma;
	}

	protected double getDelta() {
		return delta;
	}
	
	
	public void enableS9RSModifications()
	{
		// S9 mod
		delta = 0.9;
		
		// RS mod
		stateMap.put("SHRINK_3", stateMap.get("SHRINK_2")); // move SHRINK_2 to SHRINK_3
		State SHRINK_2_NEW = new State("SHRINK_2") { // new SHRINK_2
			double[] xs0;
			long prevSeed;
			
			@Override
			public double[] getX() {
//				prevSeed = seed;
//				seed = System.currentTimeMillis();
				xs0 = points.get(0).x;
				return xs0;
			}
			@Override
			public void modelComputed(double value) {
				path.append(".S0");
				double fs0 = value;
				points.get(0).x = xs0;
				points.get(0).fx = fs0;
//				seed = prevSeed;
				goTo("SHRINK_3"); // old SHRINK_2
			}
		};
	}
}
